﻿// COPYRIGHT 2012, Justin Michalczak
//
// This file is licenced to the public under the Apache
// License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License.
// You may obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using CopperVineSoftware.Stackr.Parser.Lexemes;

namespace CopperVineSoftware.Stackr.Compiler.DCPU16Assembly
{
    /// <summary>Compiles statements into DCPU-16 assembly.</summary>
    class StatementCompiler
    {
        private readonly IDictionary<string, Definition> nameTable;
        private string blockName;

        /// <summary>Sets the current block name for the compiler.</summary>
        public string BlockName
        {
            set
            {
                this.blockName = value;
            }
        }

        /// <summary>Constructor.</summary>
        /// <param name="nameTable">Table mapping names to definitions.</param>
        /// <exception cref="ArgumentNullException">If <paramref name="nameTable"/> is <c>null</c>.</exception>
        public StatementCompiler(IDictionary<string, Definition> nameTable)
        {
            if (nameTable == null)
                throw new ArgumentNullException("nameTable");

            this.nameTable = nameTable;
        }

        /// <summary>Constructor.</summary>
        /// <param name="definitions">Definitions to use.</param>
        /// <exception cref="ArgumentNullException">If <paramref name="definitions"/> is <c>null</c>.</exception>
        public StatementCompiler(IEnumerable<Definition> definitions)
        {
            if (definitions == null)
                throw new ArgumentNullException("definitions");

            this.nameTable = new Dictionary<string, Definition>();
            foreach (Definition definition in definitions)
                this.nameTable.Add(definition.Name, definition);
        }

        /// <summary>Compiles a sequence of statements.</summary>
        /// <param name="statements">Statements to compile.</param>
        /// <returns>The compiled instructions for the statements.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="statements"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(IEnumerable<Statement> statements)
        {
            if (statements == null)
                throw new ArgumentNullException("statements");

            IEnumerable<string> instructions = Enumerable.Empty<string>();
            foreach (Statement statement in statements)
                instructions = instructions.Concat(this.Compile(statement));

            return instructions;
        }

        /// <summary>Compiles a statement.</summary>
        /// <param name="statement">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="statement"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(Statement statement)
        {
            if (statement == null)
                throw new ArgumentNullException("statement");

            if (statement is IntegerConstant)
                return this.Compile(statement as IntegerConstant);
            else if (statement is NamedValue)
                return this.Compile(statement as NamedValue);
            else if (statement is EqualsConditional)
                return this.Compile(statement as EqualsConditional);
            else if (statement is NotEqualsConditional)
                return this.Compile(statement as NotEqualsConditional);
            else if (statement is GreaterThanConditional)
                return this.Compile(statement as GreaterThanConditional);
            else if (statement is LessThanConditional)
                return this.Compile(statement as LessThanConditional);
            else if (statement is WhileEquals)
                return this.Compile(statement as WhileEquals);
            else if (statement is WhileNotEquals)
                return this.Compile(statement as WhileNotEquals);
            else if (statement is WhileGreaterThan)
                return this.Compile(statement as WhileGreaterThan);
            else if (statement is WhileLessThan)
                return this.Compile(statement as WhileLessThan);
            else if (statement is Times)
                return this.Compile(statement as Times);
            else
                throw new NotImplementedException(string.Format("Statement type ({0}) not implemented.", statement.GetType().Name));
        }

        /// <summary>Compiles an integer constant statement.</summary>
        /// <param name="integerConstant">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="integerConstant"/> is <c>null</c>.</exception>
        public IEnumerable<String> Compile(IntegerConstant integerConstant)
        {
            if (integerConstant == null)
                throw new ArgumentNullException("integerConstant");

            return Push(integerConstant.Value);
        }

        /// <summary>Compiles a named value statement.</summary>
        /// <param name="namedValue">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="namedValue"/> is <c>null</c>.</exception>
        public IEnumerable<String> Compile(NamedValue namedValue)
        {
            if (namedValue == null)
                throw new ArgumentNullException("namedValue");

            if (!(nameTable.ContainsKey(namedValue.Name)))
                throw new KeyNotFoundException(string.Format("Name ({0}) undefined.", namedValue.Name));

            Definition definition = this.nameTable[namedValue.Name];
            if (definition is IntegerConstantDefinition)
            {
                IntegerConstantDefinition integerConstantDefinition = definition as IntegerConstantDefinition;
                return Push(integerConstantDefinition.Value);
            }
            else if (definition is FunctionDefinition)
                return ExecuteFunction(definition.Name);
            else if (definition is BuiltInFunctionDefinition)
                return ExecuteBuiltInFunction(definition.Name);
            else
                throw new NotImplementedException(string.Format("Definition type ({0}) not implemented.", definition.GetType().Name));
        }

        public IEnumerable<string> Compile(EqualsConditional equalsConditional)
        {
            if (equalsConditional == null)
                throw new ArgumentNullException("equalsConditional");

            string originalBlockName = this.blockName;
            string trueBlockLabel = string.Format("{0}_t", this.blockName);
            string falseBlockLabel = string.Format("{0}_f", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "SET A, POP";
            yield return "IFN A, PEEK";
            yield return string.Format("  SET PC, {0}", falseBlockLabel);

            // Compile true branch.
            this.blockName = trueBlockLabel;
            yield return string.Format(":{0}", trueBlockLabel);
            foreach (string instruction in this.Compile(equalsConditional.TrueBranch))
                yield return instruction;
            yield return string.Format("SET PC, {0}", endBlockLabel);

            // Compile false branch.
            this.blockName = falseBlockLabel;
            yield return string.Format(":{0}", falseBlockLabel);
            foreach (string instruction in this.Compile(equalsConditional.FalseBranch))
                yield return instruction;

            // Write out end label.
            yield return string.Format(":{0}", endBlockLabel);

            this.blockName = originalBlockName;
        }

        public IEnumerable<string> Compile(NotEqualsConditional notEqualsConditional)
        {
            if (notEqualsConditional == null)
                throw new ArgumentNullException("notEqualsConditional");

            string originalBlockName = this.blockName;
            string trueBlockLabel = string.Format("{0}_t", this.blockName);
            string falseBlockLabel = string.Format("{0}_f", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "SET A, POP";
            yield return "IFE A, PEEK";
            yield return string.Format("  SET PC, {0}", falseBlockLabel);

            // Compile true branch.
            this.blockName = trueBlockLabel;
            yield return string.Format(":{0}", trueBlockLabel);
            foreach (string instruction in this.Compile(notEqualsConditional.TrueBranch))
                yield return instruction;
            yield return string.Format("SET PC, {0}", endBlockLabel);

            // Compile false branch.
            this.blockName = falseBlockLabel;
            yield return string.Format(":{0}", falseBlockLabel);
            foreach (string instruction in this.Compile(notEqualsConditional.FalseBranch))
                yield return instruction;

            // Write out end label.
            yield return string.Format(":{0}", endBlockLabel);

            this.blockName = originalBlockName;
        }

        public IEnumerable<string> Compile(GreaterThanConditional greaterThanConditional)
        {
            if (greaterThanConditional == null)
                throw new ArgumentNullException("greaterThanConditional");

            string originalBlockName = this.blockName;
            string trueBlockLabel = string.Format("{0}_t", this.blockName);
            string falseBlockLabel = string.Format("{0}_f", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "SET A, POP";
            yield return "IFG PEEK, A";
            yield return string.Format("  SET PC, {0}", falseBlockLabel);

            // Compile true branch.
            this.blockName = trueBlockLabel;
            yield return string.Format(":{0}", trueBlockLabel);
            foreach (string instruction in this.Compile(greaterThanConditional.TrueBranch))
                yield return instruction;
            yield return string.Format("SET PC, {0}", endBlockLabel);

            // Compile false branch.
            this.blockName = falseBlockLabel;
            yield return string.Format(":{0}", falseBlockLabel);
            foreach (string instruction in this.Compile(greaterThanConditional.FalseBranch))
                yield return instruction;

            // Write out end label.
            yield return string.Format(":{0}", endBlockLabel);

            this.blockName = originalBlockName;
        }

        public IEnumerable<string> Compile(LessThanConditional lessThanConditional)
        {
            if (lessThanConditional == null)
                throw new ArgumentNullException("lessThanConditional");

            string originalBlockName = this.blockName;
            string trueBlockLabel = string.Format("{0}_t", this.blockName);
            string falseBlockLabel = string.Format("{0}_f", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "SET A, POP";
            yield return "IFG A, PEEK";
            yield return string.Format("  SET PC, {0}", falseBlockLabel);

            // Compile true branch.
            this.blockName = trueBlockLabel;
            yield return string.Format(":{0}", trueBlockLabel);
            foreach (string instruction in this.Compile(lessThanConditional.TrueBranch))
                yield return instruction;
            yield return string.Format("SET PC, {0}", endBlockLabel);

            // Compile false branch.
            this.blockName = falseBlockLabel;
            yield return string.Format(":{0}", falseBlockLabel);
            foreach (string instruction in this.Compile(lessThanConditional.FalseBranch))
                yield return instruction;

            // Write out end label.
            yield return string.Format(":{0}", endBlockLabel);

            this.blockName = originalBlockName;
        }

        /// <summary>Compiles a while-equals statement.</summary>
        /// <param name="whileEquals">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="whileEquals"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(WhileEquals whileEquals)
        {
            if (whileEquals == null)
                throw new ArgumentNullException("whileEquals");

            string originalBlockName = this.blockName;
            string loopBlockLabel = string.Format("{0}_l", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "ADD J, 1";
            yield return "SET [J], POP";
            yield return string.Format(":{0}", loopBlockLabel);
            yield return "IFN [J], PEEK";
            yield return string.Format("  SET PC, {0}", endBlockLabel);

            // Compile loop body
            this.blockName = loopBlockLabel;
            foreach (string instruction in this.Compile(whileEquals.Body))
                yield return instruction;

            yield return string.Format("SET PC, {0}", loopBlockLabel);
            yield return string.Format(":{0}", endBlockLabel);
            yield return "SUB J, 1";

            this.blockName = originalBlockName;
        }

        /// <summary>Compiles a while-not-equals statement.</summary>
        /// <param name="whileNotEquals">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="whileNotEquals"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(WhileNotEquals whileNotEquals)
        {
            if (whileNotEquals == null)
                throw new ArgumentNullException("whileNotEquals");

            string originalBlockName = this.blockName;
            string loopBlockLabel = string.Format("{0}_l", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "ADD J, 1";
            yield return "SET [J], POP";
            yield return string.Format(":{0}", loopBlockLabel);
            yield return "IFE [J], PEEK";
            yield return string.Format("  SET PC, {0}", endBlockLabel);

            // Compile loop body
            this.blockName = loopBlockLabel;
            foreach (string instruction in this.Compile(whileNotEquals.Body))
                yield return instruction;

            yield return string.Format("SET PC, {0}", loopBlockLabel);
            yield return string.Format(":{0}", endBlockLabel);
            yield return "SUB J, 1";

            this.blockName = originalBlockName;
        }

        /// <summary>Compiles a while-greater-than statement.</summary>
        /// <param name="whileGreaterThan">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="whileGreaterThan"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(WhileGreaterThan whileGreaterThan)
        {
            if (whileGreaterThan == null)
                throw new ArgumentNullException("whileGreaterThan");

            string originalBlockName = this.blockName;
            string loopBlockLabel = string.Format("{0}_l", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "ADD J, 1";
            yield return "SET [J], POP";
            yield return string.Format(":{0}", loopBlockLabel);
            yield return "IFG [J], PEEK";
            yield return string.Format("  SET PC, {0}", endBlockLabel);
            yield return "IFE [J], PEEK";
            yield return string.Format("  SET PC, {0}", endBlockLabel);

            // Compile loop body
            this.blockName = loopBlockLabel;
            foreach (string instruction in this.Compile(whileGreaterThan.Body))
                yield return instruction;

            yield return string.Format("SET PC, {0}", loopBlockLabel);
            yield return string.Format(":{0}", endBlockLabel);
            yield return "SUB J, 1";

            this.blockName = originalBlockName;
        }

        /// <summary>Compiles a while-greater-than statement.</summary>
        /// <param name="whileLessThan">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="whileLessThan"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(WhileLessThan whileLessThan)
        {
            if (whileLessThan == null)
                throw new ArgumentNullException("whileLessThan");

            string originalBlockName = this.blockName;
            string loopBlockLabel = string.Format("{0}_l", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "ADD J, 1";
            yield return "SET [J], POP";
            yield return string.Format(":{0}", loopBlockLabel);
            yield return "IFG [J], PEEK";
            yield return string.Format("  SET PC, {0}", endBlockLabel);
            yield return "IFE [J], PEEK";
            yield return string.Format("  SET PC, {0}", endBlockLabel);

            // Compile loop body
            this.blockName = loopBlockLabel;
            foreach (string instruction in this.Compile(whileLessThan.Body))
                yield return instruction;

            yield return string.Format("SET PC, {0}", loopBlockLabel);
            yield return string.Format(":{0}", endBlockLabel);
            yield return "SUB J, 1";

            this.blockName = originalBlockName;
        }

        /// <summary>Compiles a times statement.</summary>
        /// <param name="times">Statement to compile.</param>
        /// <returns>The compiled instructions for the statement.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="times"/> is <c>null</c>.</exception>
        public IEnumerable<string> Compile(Times times)
        {
            if (times == null)
                throw new ArgumentNullException("times");

            string originalBlockName = this.blockName;
            string loopBlockLabel = string.Format("{0}_l", this.blockName);
            string endBlockLabel = string.Format("{0}_e", this.blockName);

            yield return "ADD J, 1";
            yield return "SET [J], POP";
            yield return string.Format(":{0}", loopBlockLabel);
            yield return "IFE [J], 0";
            yield return string.Format("  SET PC, {0}", endBlockLabel);

            // Compile loop body
            this.blockName = loopBlockLabel;
            foreach (string instruction in this.Compile(times.Body))
                yield return instruction;

            yield return "SUB [J], 1";
            yield return string.Format("SET PC, {0}", loopBlockLabel);
            yield return string.Format(":{0}", endBlockLabel);
            yield return "SUB J, 1";

            this.blockName = originalBlockName;
        }

        /// <summary>Creates instructions for a push operation.</summary>
        /// <param name="value">Value to push.</param>
        /// <returns>Instructions pushing the given value to the stack.</returns>
        private static IEnumerable<string> Push(ushort value)
        {
            yield return string.Format("SET PUSH, {0}", value);
        }

        /// <summary>Creates instructions for a function call.</summary>
        /// <param name="functionName">Name of function to call.</param>
        /// <returns>Instructions calling the given function.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="functionName"/> is <c>null</c>.</exception>
        private static IEnumerable<string> ExecuteFunction(string functionName)
        {
            if (string.IsNullOrEmpty(functionName))
                throw new ArgumentNullException("functionName");

            yield return string.Format("SET PUSH, {0}", functionName);
            yield return string.Format("JSR exec");
        }

        /// <summary>Creates instructions for a built-in function call.</summary>
        /// <param name="functionName">Name of function to call.</param>
        /// <returns>Instructions calling the given built-in function.</returns>
        /// <exception cref="ArgumentNullException">If <paramref name="functionName"/> is <c>null</c>.</exception>
        private static IEnumerable<string> ExecuteBuiltInFunction(string functionName)
        {
            yield return string.Format("JSR {0}", functionName);
        }
    }
}
